# koroFileHeader at Yungoal acer
# Create: 2021-12-31 10:06:08
# LastEdit: 2022-01-04 13:13:48
"""Azure Windows 补丁更新
Run Command 接口的问题，执行完毕后才能获取到全部的执行结果，无法入命令行中执行那样，实时查看输出。
Run Command 最长允许脚本执行90分钟，如果执行时间过长，就可能拿不到脚本返回。
如果脚本自带重启系统，那么也无法获取到中间的返回信息，所以不考虑完成后自动重启。
要想如命令行中那样实时获取到结果，可以考虑引入日志系统，将输出发送到日志系统。
"""
__author__ = '749B'

import os, sys
BASE_DIR = os.getcwd()
sys.path.append(BASE_DIR)

import argparse
import textwrap
import time
import msrestazure.azure_cloud
# pylint: disable = import-error
from azure_client import Azure_Client
from azure_client import data_format
from azure_client.models import RuncmdAsyncJob
from utils.tools import remove_empty
from utils.blueking import job_start, job_success, job_fail
from utils.env import init_environs

from typing import Optional, List
# 只是因为函数的返回需要做类型判断
from azure.mgmt.resourcegraph.models._models_py3 import QueryResponse


def str_query_base(
    sub_id: Optional[str] = None,
    rg: Optional[str] = None,
    filters: Optional[str] = None,
    hosts: Optional[List[str]] = None,
    ) -> str:
    """包含过滤条件的查询语句，在这个基础上过滤虚拟机"""
    str_query = textwrap.dedent("""\
    Resources
    | where type == "microsoft.compute/virtualmachines"
    | extend 
        OSType = properties.storageProfile.osDisk.osType,
        nicid = tostring(properties.networkProfile.networkInterfaces[0].id)
    | where OSType =~ "windows"
    {filter_sub_id}
    {filter_rg}
    {other_filters}
    | project resourceGroup, OSType, subscriptionId, vm_name=name, nicid
    | join kind=leftouter (
        Resources
        | where type =~ "microsoft.network/networkinterfaces"
        | project nicid=id, ip_address=properties.ipConfigurations[0].properties.privateIPAddress
    ) on nicid 
    | project-away nicid, nicid1
    | join kind=leftouter (
        Resourcecontainers
        | where type =~ "microsoft.resources/subscriptions"
        | project SubName=name, subscriptionId
    ) on subscriptionId
    | project-away subscriptionId1
    {filter_hosts}
    """).format(
        filter_sub_id=f'| where subscriptionId =~ "{sub_id}"' if sub_id else '',
        filter_rg=f'| where resourceGroup =~ "{rg}"' if rg else '',
        other_filters=f'| where {filters}' if filters else '',
        filter_hosts='| where vm_name in ({0}) or ip_address in ({0})'.format(
            ','.join(['"' + i + '"' for i in hosts])
            ) if hosts else '',
        )
    # 去掉空白行
    str_query = remove_empty(str_query)
    return str_query

def get_vm_count(client: Azure_Client, base_query: str) -> int:
    """返回虚拟机的数量"""
    strQuery = base_query + '| summarize count()\n'
    # 去掉空白行
    strQuery = remove_empty(strQuery)
    res = client.resource_graph_query_default(
        strQuery,
        [sub.get("subscription_id") for sub in client.get_subscriptions()],
    )
    return res.data[0].get('count_')

def get_vm_info(
    client: Azure_Client,
    base_query: str,
    limit: Optional[int] = None,
    ) -> QueryResponse:
    """拉取所有的虚拟机"""
    strQuery = base_query + f'| limit {limit}\n' if limit else base_query 
    # 去掉空白行
    strQuery = remove_empty(strQuery)
    res = client.resource_graph_query_default(
        strQuery,
        [sub.get("subscription_id") for sub in client.get_subscriptions()],
    )
    return res

def init_argparse(args: Optional[List[str]] = None) -> argparse.Namespace:
    """初始化: 命令行参数"""
    parser = argparse.ArgumentParser(description='Windows补丁更新')

    parser.add_argument('hosts', nargs='*', help='虚拟机的名称或第一个内网IP')
    parser.add_argument('-e', '--envfile', help="存有用户认证信息的env文件的文件名")
    parser.add_argument('-s', '--subscription', help='限定订阅，订阅名称或订阅ID都可')
    parser.add_argument('-g', '--resource_group', help='限定资源组名称')
    parser.add_argument('-f', '--filters', help='其他过滤条件，参考资源查询的where语句')
    parser.add_argument('--check_only', action='store_true', help='仅检查，不更新')
    parser.add_argument('--test_mode', default='off', choices=['off', 'vm_info', 'runcmd_check'], help='测试模式')
    parser.add_argument('-t', '--timeout', type=int, default=300, help='超过多少时间后，不再继续等待命令结果')
    parser.add_argument('--pool', type=int, default=50, help='同时进行的任务的数量')
    parser.add_argument('--show_query', action='store_true', help='只打印KQL查询语句，检查查询语句问题')
    parser.add_argument('--limit', type=int, dest='dev_limit', help='调试用，过滤前几条记录执行任务')
    parser.add_argument('--show_count', action='store_true', help='每行记录去打印序号')
    parser.add_argument('--csvfile', dest='dev_csvfile', help="仅限本地使用，输出一份到文件")

    args = parser.parse_args(args)
    return args

def main(args: Optional[List[str]] = None) -> None:
    """主函数"""
    # 初始化命令行参数
    if args is None:
        # 蓝鲸作业平台有BUG，重做的任务，空着的参数里都会被填充null
        args = [arg for arg in sys.argv[1:] if arg != 'null' and not arg.endswith('=null')]
    args = init_argparse(args)
    if not args.envfile:
        return job_fail("请提供存有用户认证信息的env文件的文件名")

    # 初始化环境变量
    envs = init_environs(args.envfile)
    if not envs:
        return job_fail("env文件解密错误")
    if envs['runtime_env'] == 'blueking':
        for key in args.__dict__.keys():
            if key.startswith('dev_'):
                v = getattr(args, key)
                if v:
                    print(f'{key}={v}: 不支持的命令行参数，已经忽略。', file=sys.stderr)
                    setattr(args, key, None)

    # 实例化客户端
    client = Azure_Client(getattr(msrestazure.azure_cloud, envs['cloud_name']))

    if args.subscription:
        # pylint: disable = no-member
        subscription = client.get_subscriptionid_by_name(args.subscription)
        if not subscription:
            return job_fail("订阅或订阅ID不存在:%s" % args.subscription)
        args.subscription = subscription
    base_query = str_query_base(args.subscription, args.resource_group, args.filters, args.hosts)
    if args.show_query:
        print(base_query)  # 蓝鲸平台上，那么多内容放到下面的job_success里会返回错误。单独print出来。
        return job_success("已打印查询语句，程序结束。")

    # 检查虚拟机数量
    vm_count = get_vm_count(client, base_query)
    print("虚拟机数量:", vm_count)
    if not vm_count:
        return job_fail("未匹配到任何虚拟机，建议检查过滤语句。")

    # 查询虚拟机信息
    vm_info = get_vm_info(client, base_query, args.dev_limit)
    if args.test_mode == 'vm_info':
        data_format.output_csv(vm_info.data, comma_begin=False)
        return job_success("打印匹配的VM信息，程序结束。")

    # 拼接脚本
    script = None
    if script is None and args.test_mode == 'runcmd_check':
        script = SCRIPT_RUNCMD_CHECK
    if script is None:
        script = SCRIPT_CHECK
        if not args.check_only:
            script.extend(SCRIPT_UPDATA)

    # 异步执行任务
    job = RuncmdAsyncJob(client, args.pool, args.timeout, show_count=args.show_count, multline=True, csvfile=args.dev_csvfile)
    for vm in vm_info.data:
        # vm 是个字典，下面是全部的字段
        # vm_name = vm.get('vm_name')
        # vm_resource_group = vm.get('resourceGroup')
        # vm_ostype = vm.get('OSType')
        # vm_subscription_id = vm.get('subscriptionId')
        # vm_subscription_name = vm.get('SubName')
        # vm_ip_address = vm.get('ip_address')
        pool_is_full = job.add_task(
            vm, ['step', 'msg'], script,
            parameters=[{'name': 'check_only', 'value': 1 if args.check_only else ''}])
        if pool_is_full:
            job.get_task_result()
    # 最后取完所有的任务结果
    while job.get_task_result():
        if not args.check_only:
            # 如果是更新补丁，每次轮询等待时间长一点
            time.sleep(60)
    job.close()

    # 结束，最后打印提示信息
    if args.test_mode == 'runcmd_check':
        print()
        print("注意:", "runcmd_check模式，仅运行测试命令，验证RunCommand权限和功能可用")
    elif args.check_only:
        print()
        print("注意:", "check_only开启，仅检查不更新")
    job_success()


# 检查更新的脚本
SCRIPT_CHECK = """\
param(
    [switch]$check_only
)

$needsReboot = $false
$isSucceeded = $true

# 定义更新搜索，筛选条件参考官方文档。
# https://docs.microsoft.com/en-us/windows/win32/api/wuapi/nf-wuapi-iupdatesearcher-search
$Criteria = "AutoSelectOnWebSites=1 and IsHidden=0 and IsInstalled=0"
# 检查更新
$UpdateSession = New-Object -ComObject Microsoft.Update.Session 
$UpdateSearcher = $UpdateSession.CreateUpdateSearcher()

# 搜索更新
Write-Host "Start,Searching for Updates" 
$SearchResult = $UpdateSearcher.Search($Criteria)
Write-Host "UpdateCount,Found [$($SearchResult.Updates.count)] Updates to Download and install"
if ($check_only) {
    $SearchResult.Updates | %{Write-Output "UpdateTitle,$($_.Title)"}
}
""".splitlines()
# 执行更新的脚本，需要拼接在检查更新的脚本之后一起执行
SCRIPT_UPDATA = """\
function Get-WIAStatusValue($value) { 
    switch -exact ($value) { 
        0 {"NotStarted"} 
        1 {"InProgress"} 
        2 {"Succeeded"} 
        3 {"SucceededWithErrors"} 
        4 {"Failed"} 
        5 {"Aborted"} 
    } 
}

foreach($Update in $SearchResult.Updates) {
    # Add Update to Collection 
    $UpdatesCollection = New-Object -ComObject Microsoft.Update.UpdateColl 
    if ( $Update.EulaAccepted -eq 0 ) { 
        $Update.AcceptEula() 
    } 
    $UpdatesCollection.Add($Update) | out-null 
    Write-Host "Downloading,  + Downloading Update $($Update.Title)" 
    $UpdatesDownloader = $UpdateSession.CreateUpdateDownloader() 
    $UpdatesDownloader.Updates = $UpdatesCollection 
    $DownloadResult = $UpdatesDownloader.Download() 
    $Message = "Download,    - Download {0}" -f (Get-WIAStatusValue $DownloadResult.ResultCode) 
    Write-Host $message
    if ( $DownloadResult.ResultCode -ne 2 ) {
        $isSucceeded = $false
        continue
    }
    Write-Host "Installing,    - Installing Update" 
    $UpdatesInstaller = $UpdateSession.CreateUpdateInstaller() 
    $UpdatesInstaller.Updates = $UpdatesCollection 
    $InstallResult = $UpdatesInstaller.Install() 
    $Message = "Install,    - Install {0}" -f (Get-WIAStatusValue $InstallResult.ResultCode) 
    Write-Host $message 
    if ( $InstallResult.ResultCode -ne 2 ) {
        $isSucceeded = $false
        continue
    }
    if ( $installResult.rebootRequired ) {
        $needsReboot = $true
    }
}

if ( $isSucceeded ) {
    if ( $needsReboot ) {
        Write-Host "Succeeded,Succeeded but Needs Reboot"
        # 下面是直接重启机器，注释掉了
        # restart-computer
    } else {
        Write-Host "Succeeded,Succeeded and Reboot not required :-)"
    }
} else {
    Write-Host "Problem,Check the Problem :-("
    exit 1
}
""".splitlines()
# 测试Run Command的简单脚本
SCRIPT_RUNCMD_CHECK = """\
Write-Output "hostname,$(hostname)"
Write-Output "date,$(date)"
""".splitlines()

    
if __name__ == '__main__':

    job_start()

    ###### 可在此处开始编写您的脚本逻辑代码
    ###### iJobs中执行脚本成功和失败的标准只取决于脚本最后一条执行语句的返回值
    ###### 如果返回值为0，则认为此脚本执行成功，如果非0，则认为脚本执行失败

    main()
